import json
from typing import Dict, List, Optional
from mem0.configs.llms.base import BaseLlmConfig
import random
from http import HTTPStatus
from dashscope import Generation
from mem0.llms.base import LLMBase


class QianWenLLM(LLMBase):
    def __init__(self, config: Optional[BaseLlmConfig] = None):
        super().__init__(config)

        if not self.config.model:
            self.config.model = "qwen-max"

    def _parse_response(self, response, tools):
        """
        Process the response based on whether tools are used or not.

        Args:
            response: The raw response from API.
            tools: The list of tools provided in the request.

        Returns:
            str or dict: The processed response.
        """
        if tools:
            processed_response = {
                "content": response['output']['choices'][0]['message']['content'],
                "tool_calls": []
            }

            # Check if there are any tool calls in the response
            if 'tool_calls' in response['output']['choices'][0]['message'] and response['output']['choices'][0]['message'].tool_calls:
                for tool_call in response['output']['choices'][0]['message'].tool_calls:
                    processed_response["tool_calls"].append({
                        "name": tool_call['function']['name'],
                        "arguments": json.loads(tool_call['function']['arguments'])
                    })

            return processed_response
        else:
            return response['output']['choices'][0]['message']['content']

    def generate_response(
            self,
            messages: List[Dict[str, str]],
            result_format=None,
            tools: Optional[List[Dict]] = None,
            tool_choice: str = "auto",
    ):

        params = {
            "model": self.config.model,
            "messages": messages,
            "seed": random.randint(1, 10000),
            "temperature": 0.8,
            "top_p": 0.8,
            "top_k": 50,
            "result_format": 'message'
        }
        if result_format:
            params["result_format"] = result_format

        if tools is not None:
            params["tools"] = tools
            params["tool_choice"] = tool_choice

        response = Generation.call(**params)

        if response.status_code == HTTPStatus.OK:
            parsed_response = self._parse_response(response, tools)
            return parsed_response
        else:
            raise Exception(f'Request id: {response.request_id}, '
                            f'Status code: {response.status_code}, '
                            f'error code: {response.code}, '
                            f'error message: {response.message}')